{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mindspore import nn\n",
    "from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits\n",
    "from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor\n",
    "from mindspore.train import Model\n",
    "from mindspore.nn.metrics import Accuracy\n",
    "import mindspore.dataset as mds\n",
    "from mindspore import context\n",
    "from mindspore import Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"CPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LR(nn.Cell):\n",
    "    def __init__(self, inputs_size=2, outputs_size=2):\n",
    "        super(LR, self).__init__()\n",
    "        self.layer_1 = nn.Dense(inputs_size, 8, activation='relu')\n",
    "        self.layer_ouput = nn.Dense(8, outputs_size, activation='relu')\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.layer_1(x)\n",
    "        x = self.layer_ouput(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model():\n",
    "    network = LR(2, 2)\n",
    "    net_loss = SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True)\n",
    "    net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)\n",
    "    model = Model(network, net_loss, net_opt, metrics={\"Accuracy\": Accuracy()})\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_generator_func(dataset_size):\n",
    "    def generator_func():\n",
    "        for i in range(dataset_size):\n",
    "            X = np.random.rand(2).astype('float32')\n",
    "            y = X[0] > X[1]\n",
    "            y = np.array([y]).astype('int32')\n",
    "            yield (X, y)\n",
    "    return generator_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_datasets():\n",
    "    ds1 = mds.GeneratorDataset(get_generator_func(100000), [\"feature\", \"label\"])\n",
    "    ds1.set_dataset_size(32)\n",
    "    ds2 = ds1.shuffle(buffer_size=5)\n",
    "    ds3 = ds2.batch(32, drop_remainder=True)\n",
    "    return ds3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    datasets = create_datasets()\n",
    "    model = build_model()\n",
    "    config_ck = CheckpointConfig(save_checkpoint_steps=10, keep_checkpoint_max=10)\n",
    "    checkpoint_callback = ModelCheckpoint(prefix=\"checkpoint_LR\", directory='./model/lr', config=config_ck)\n",
    "    model.train(10, datasets, callbacks=[checkpoint_callback, LossMonitor()], dataset_sink_mode=False)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "epoch: 1 step: 1, loss is 0.6931441\nepoch: 2 step: 1, loss is 0.69308287\nepoch: 3 step: 1, loss is 0.6932115\nepoch: 4 step: 1, loss is 0.69303876\nepoch: 5 step: 1, loss is 0.6934409\nepoch: 6 step: 1, loss is 0.69381046\nepoch: 7 step: 1, loss is 0.6940061\nepoch: 8 step: 1, loss is 0.69308174\nepoch: 9 step: 1, loss is 0.6931391\nepoch: 10 step: 1, loss is 0.6931475\n"
    }
   ],
   "source": [
    "model = train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore.train.serialization import load_checkpoint, load_param_into_net, export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_dict = load_checkpoint(ckpoint_file_name='./model/lr/checkpoint_LR-10_1.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_model = LR(2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_param_into_net(save_model, param_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "origin_data = np.random.uniform(low=0, high=255, size=(32, 2)).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = Tensor(origin_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "export(save_model, input_data, file_name='./lr.ms', file_format='LITE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "hello\n"
    }
   ],
   "source": [
    "print(\"hello\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python37164bit961d820658604e08a9cb598b3849190f",
   "display_name": "Python 3.7.1 64-bit"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}